import json
import os
import re

import requests
from PyPDF2 import PdfReader

LLM_HOST = "http://10.100.8.72:18888/v1/chat/completions"


def send_chat_request(content):
    for attempt in range(3):  # 尝试最多3次
        chat_req = {
            "model": "openchat_3.5",
            "temperature": "0.1",
            "messages": [
                {
                    "role": "user",
                    "content": f"需提取内容：{content}\r\n\r\n。请提取出文件编号,以fileNo返回json格式，如果提取失败，请返回空字符串，例子(编号：ZXLH01QD180807，编号为T2020093000009)，并以json格式返回，只输出json'"
                }
            ]
        }

        try:
            response = requests.post(LLM_HOST, json=chat_req)
            response.raise_for_status()
            wrapper = response.json()

            if 'choices' in wrapper and len(wrapper['choices']) > 0:
                file_details_content = wrapper['choices'][0]['message']['content']
                try:
                    file_details = json.loads(file_details_content)
                    # 检查是否包含fileNo键
                    if 'fileNo' in file_details:
                        return file_details, None
                except json.JSONDecodeError:
                    # 如果结果不是json格式，继续尝试
                    continue
            else:
                raise ValueError("wrapper['choices'] is empty or missing")

        except requests.RequestException as e:
            return None, e

    # 如果3次尝试后仍无法获取有效的json或fileNo，返回错误
    return None, ValueError("Failed to retrieve valid json with 'fileNo' after 3 attempts")


def find_and_write_matches(pdf_path, output_file):
    try:
        print("正在处理：", pdf_path)
        # dir_path, file_name = os.path.split(pdf_path)
        # pdf_path = os.path.join(dir_path, '3943416663323578368.pdf')

        reader = PdfReader(pdf_path)
        all_text = ""
        for page in reader.pages:
            text = page.extract_text()
            if text:
                all_text += text
                if len(all_text) > 4000:
                    all_text = all_text[:4000]
                    break

        # 处理合并后的文本
        if all_text:
            file_details, error = send_chat_request(all_text)
            file_name_without_extension = os.path.splitext(
                os.path.basename(pdf_path))[0]
            if file_details:
                output_file.write(
                    f"{file_name_without_extension}={file_details['fileNo'].replace(' ', '')}\n")
                output_file.flush()
            else:
                output_file.write(
                    f"{file_name_without_extension}={''}\n")
                output_file.flush()
                print("出错了：", error)
        else:
            print(f"No text found in {os.path.basename(pdf_path)}")

        print(f"Processed {os.path.basename(pdf_path)}")
    except Exception as e:
        print(f"Failed to process {os.path.basename(pdf_path)}: {e}")


def convert_pdf_folder_to_single_txt(pdf_folder, output_txt_path):
    with open(output_txt_path, 'w', encoding='utf-8') as output_file:
        for filename in os.listdir(pdf_folder):
            if filename.lower().endswith('.pdf'):
                pdf_path = os.path.join(pdf_folder, filename)
                find_and_write_matches(pdf_path, output_file)


def convert_pdf_readtxt_to_single_txt(pdf_folder, pdf_txt_path, output_txt_path):
    with open(output_txt_path, 'w', encoding='utf-8') as output_file:
        with open(pdf_txt_path, 'r', encoding='utf-8') as txt_list_file:
            for line in txt_list_file:
                filename_without_extension = line.strip()  # 移除行尾的换行符
                pdf_filename = f"{filename_without_extension}.pdf"  # 添加.pdf后缀
                pdf_path = os.path.join(pdf_folder, pdf_filename)
                if os.path.isfile(pdf_path):
                    find_and_write_matches(pdf_path, output_file)
                else:
                    print(f"File {pdf_path} does not exist.")


if __name__ == '__main__':
    # The path to the PDF folder is one level up from the current script directory
    pdf_folder_path = os.path.join(os.path.dirname(__file__), '..', 'pdf')
    pdf_txt_path = os.path.join(os.path.dirname(
        __file__), '..', 'filenames_fileid.txt')
    pdf_folder_path = os.path.abspath(pdf_folder_path)  # Normalize the path
    pdf_txt_path = os.path.abspath(pdf_txt_path)  # Normalize the path

    # Define the output TXT file path
    output_txt_filename = 'filenames.txt'
    output_txt_path = os.path.join(
        os.path.dirname(__file__), '..', output_txt_filename)

    # convert_pdf_folder_to_single_txt(pdf_folder_path, output_txt_path)
    convert_pdf_readtxt_to_single_txt(
        pdf_folder_path, pdf_txt_path, output_txt_path)
